import sys, os
sys.path.append(os.pardir)
import numpy as np
import matplotlib.pyplot as plt
from dataset.mnist import load_mnist
from twoLayerNet import TwoLayerNet

# 读入数据
(x_train, t_train), (x_test, t_test) = load_mnist(normalize=True, one_hot_label=True)

# 初始化神经网络，由于我们输入的图像为28*28标准，所以输入层神经元数量应该为784，隐藏层我们设置一个合适数值50，输出层呢，类别为十个，所以设置为10
network = TwoLayerNet(input_size=784, hidden_size=50, output_size=10)

iters_num = 10000  # 适当设定循环次数
train_size = x_train.shape[0]
batch_size = 100  # 每次抽取100份来学习
learning_rate = 0.1  # 学习率

train_loss_list = []  # 用来记录学习过程中损失函数值的变化吧
train_acc_list = []  # 用来记录学习过程中精确度的变化吧
test_acc_list = []  # 用来记录测试的精确度？


iter_per_epoch = max(train_size / batch_size, 1)

for i in range(iters_num):
    # print(str(i)+ "  ",end='')
    batch_mask = np.random.choice(train_size, batch_size)  # 从总的训练数据中随机抽取固定数量的数据（这里是以数据索引代表数据）
    x_batch = x_train[batch_mask]  # 抽取好的训练数据
    t_batch = t_train[batch_mask]  # 与之对应的正确标签

    # 计算梯度
    # grad = network.numerical_gradient(x_batch, t_batch)  # 低速版，实在等不动了
    grad = network.gradient(x_batch, t_batch)  # 计算梯度，高速版！！

    # 更新参数（使用梯度下降法）
    for key in ('W1', 'b1', 'W2', 'b2'):
        network.params[key] -= learning_rate * grad[key]

    loss = network.loss(x_batch, t_batch)  # 梯度法调参后，再查看一下损失函数值
    train_loss_list.append(loss)  # 将其记录下来

    # 怎么，这里是认为学习完一遍所有数据之后，就完整测一遍所有数据是吧
    # 一个单位的epoch之后，要测试一下学习成果
    if i % iter_per_epoch == 0:
        train_acc = network.accuracy(x_train, t_train)
        test_acc = network.accuracy(x_test, t_test)
        train_acc_list.append(train_acc)
        test_acc_list.append(test_acc)
        print("train acc, test acc | " + str(train_acc) + "," + str(test_acc))

# 绘制图形
markers = {'train': 'o', 'test': 's'}
x = np.arange(len(train_acc_list))
plt.plot(x, train_acc_list, label='train ass')
plt.plot(x, test_acc_list, label='test ass', linestyle='--')
plt.xlabel("epochs")
plt.ylabel("accuracy")
plt.ylim(0, 1.0)
plt.legend(loc='lower right')
plt.show()
